import math
import torch.nn.functional as F
import torch


def self_attention(query, key, value, mask=None):
    d_k = query.size(-1)
    scores = torch.matmul(query, key.transpose(-2, -1) / math.sqrt(d_k))
    if mask is not None:
        mask.cuda()
        scores = scores.masked_fill(mask == 0, -1e9)
        # mask和value，将mask中取值为0位置对应于scores的相应位置用-1e9填充。
    self_attn = F.softmax(scores, dim=-1)
    return torch.matmul(self_attn, value), self_attn


print(torch.cuda.is_available())